///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/ham/ham.h
///
/// \brief   header file for the hamiltonian class
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#ifndef FORMIC_HAMILTONIAN_HEADER
#define FORMIC_HAMILTONIAN_HEADER

#include<formic/numeric/numeric.h>
#include<formic/cholesky.h>
#include<formic/lapack/interface.h>
#include<formic/ham/tei.h>
#include<formic/vector/vector.h>
#include<formic/vector/matrix.h>

namespace formic {

  class Archive;
  class InputBase;
  template<class T> class BinaryArchiveTemplate;

  void init_hamiltonian_commands();

  template<class S> void get_text_ints(const formic::InputBase &,
                                       const std::string &,
                                       double &,
                                       S * const,
                                       S * const,
                                       S * const);

  ///////////////////////////////////////////////////////////////////////////////////////////////////
  /// \brief  A class to represent the Hamiltonian
  ///
  ///////////////////////////////////////////////////////////////////////////////////////////////////
  template<class T> class Hamiltonian : public formic::CholeskySource<T>, public formic::CholeskyTarget<T> {

    private:

      /// \brief  the number of spatial orbitals
      int _norb;

      /// \brief  the nuclear repulsion and frozen core energy
      double _e_cor_nuc;

      /// \brief  the orbital coefficients relative to the S^(-1/2) atomic orbital basis
      formic::Matrix<T,int> _orb_coeff;

      /// \brief  the orbital overlap integrals
      std::vector<T> _ovi;

      /// \brief  the one electron integrals
      std::vector<T> _oei;

      /// \brief  the modified one electron integrals
      std::vector<T> _mei;

      /// \brief  the two electron integrals, stored using 8-fold symmetry
      std::vector<T> _tei;

      /// \brief  the number of cholesky vectors
      int _n_cholesky;

      /// \brief  truncation threshold for cholesky decomposition
      double _chol_thresh;

      /// \brief  an array containing the cholesky vectors
      std::vector<T> _cholesky_vecs;

      /// \brief  an array containing the transposed cholesky vectors
      std::vector<T> _cholesky_vecs_t;

      /// \brief  vector to hold the dense two electron integral array
      formic::Vector<T,int> _dense_tei;

      void initialize_mei();

      void compute_cholesky();

    public:

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief constructs an empty hamiltonian
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      Hamiltonian() {}

      Hamiltonian(const int norb);

      const formic::Matrix<T,int> & orb_coeff() const { return _orb_coeff; }

      void initialize_hubbard(const formic::InputBase & userinp);
      void initialize_integrals(const formic::InputBase & userinp,
                                const std::string & source,
                                const std::string & filename);
      void write_dmrg_ints(const formic::InputBase & userinp) const;
      void write_text_ints(const formic::InputBase & userinp, const std::string & filename) const;

      void transform_orbitals(const T * const new_orbs, const bool from_curr = true, const bool print_only = false);

      void set_integrals(const T * const new_oei, const T * const new_tei);

      void lowdin_localize();
      void lowdin_localize(const std::vector<int> & orb_inds);

      static void print_oei(const int n, const std::vector<T> & mat, const std::string & title);

      static void print_tei_diag(const int n, const std::vector<T> & mat, const std::string & title);

      /// \brief   return the nuclear repulsion and frozen core energy
      double e_cor_nuc() const { return _e_cor_nuc; }

      /// \brief   return the nuclear repulsion and frozen core energy (read/write)
      double & e_cor_nuc() { return _e_cor_nuc; }

      /// \brief   return the number of cholesky vectors
      int n_cholesky() const { return _n_cholesky; }

      /// \brief   return the cholesky vectors
      const std::vector<T> & cholesky_vecs() const { return _cholesky_vecs; }

      /// \brief   return the transposed cholesky vectors
      const std::vector<T> & cholesky_vecs_t() const { return _cholesky_vecs_t; }

      void get_pqpq_ints(std::vector<T> & pqpq_ints) const;

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief   get an overlap integral
      /// 
      /// \param[in]  p   1st index
      /// \param[in]  q   2nd index
      ///
      /// \return  the integral \f$  s_{p q}  \f$
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      T get_ovi(const int p, const int q) const { return _ovi[ p * _norb + q ]; }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief   get an overlap integral (read/write)
      /// 
      /// \param[in]  p   1st index
      /// \param[in]  q   2nd index
      ///
      /// \return  the integral \f$  s_{p q}  \f$
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      T & get_ovi(const int p, const int q) { return _ovi[ p * _norb + q ]; }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief   get a one electron integral
      /// 
      /// \param[in]  p   1st index
      /// \param[in]  q   2nd index
      ///
      /// \return  the integral \f$  h_{p q}  \f$
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      T get_oei(const int p, const int q) const { return _oei[ p * _norb + q ]; }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief   get a one electron integral (read/write)
      /// 
      /// \param[in]  p   1st index
      /// \param[in]  q   2nd index
      ///
      /// \return  the integral \f$  h_{p q}  \f$
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      T & get_oei(const int p, const int q) { return _oei[ p * _norb + q ]; }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief   get a modified one electron integral
      /// 
      ///    The modified one electron integrals f are definded as
      ///      \f$  f_{p q} = h_{p q} - \frac{1}{2} \sum_r  [pr|rq]  \f$
      ///    where h is the original one electron integrals and [**|**] are the two electron integrals.
      ///
      /// \param[in]  p   1st index
      /// \param[in]  q   2nd index
      ///
      /// \return  the modified integral \f$  f_{p q}  \f$
      ///
      /// \sa initialize_mei
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      T get_mei(const int p, const int q) const { return _mei[ p * _norb + q ]; }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief   get a two electron integral (read only)
      /// 
      /// \param[in]  p   1st index for electron coordinate 1
      /// \param[in]  q   2nd index for electron coordinate 1
      /// \param[in]  r   1st index for electron coordinate 2
      /// \param[in]  s   2nd index for electron coordinate 2
      ///
      /// \return  the integral [pq|rs]
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      T get_tei(const int p, const int q, const int r, const int s) const {
        const int ind = formic::tei_compound_index(p,q,r,s);
        assert( ind >= 0 );
        assert( ind < _tei.size() );
        return _tei[ind];
      }

      ///////////////////////////////////////////////////////////////////////////////////////////////////
      /// \brief   get a two electron integral (read/write)
      /// 
      /// \param[in]  p   1st index for electron coordinate 1
      /// \param[in]  q   2nd index for electron coordinate 1
      /// \param[in]  r   1st index for electron coordinate 2
      /// \param[in]  s   2nd index for electron coordinate 2
      ///
      /// \return  the integral [pq|rs]
      ///
      ///////////////////////////////////////////////////////////////////////////////////////////////////
      T & get_tei(const int p, const int q, const int r, const int s) {
        const int ind = formic::tei_compound_index(p,q,r,s);
        assert( ind >= 0 );
        assert( ind < _tei.size() );
        return _tei[ind];
      }

      const T * dense_oei_ptr() const { return &_oei.at(0); }

      const T * dense_tei_ptr();

      void reorder_orbitals(const formic::InputBase & userinp,
                            const std::vector<int> & order);

      /// \brief  get the dimension of the matrix to be decomposed via cholesky
      int chol_size() const { return _norb*_norb; }

      /// \brief  get the cholesky decomposition threshold
      double chol_thresh() const { return _chol_thresh; }

      /// \brief  get a diagonal element of the two electron integral matrix for use in cholesky
      double chol_diag(const int i) const {
        const int p = i / _norb;
        const int q = i % _norb;
        return formic::real(this->get_tei(p,q,p,q));
      }

      /// \brief  get a column of the two electron integral matrix for use in cholesky
      void chol_get_col(const int j, std::vector<T> & col) const {
        assert( col.size() == _norb*_norb );
        const int p = j / _norb;
        const int q = j % _norb;
        for (int r = 0; r < _norb; r++)
        for (int s = 0; s < _norb; s++)
          col[_norb*r+s] = this->get_tei(p,q,r,s);
      }

      /// \brief  prepare storage for the cholesky vectors
      void chol_prepare(const int m) {
        _n_cholesky = m;
        _cholesky_vecs.resize( _norb * _norb * m );
        _cholesky_vecs_t.resize( _norb * _norb * m );
      }

      /// \brief  store one of the cholesky vectors
      void chol_put_col(const int j, const std::vector<T> & col) {
        assert( col.size() == _norb*_norb );
        assert( j >= 0 );
        assert( j < _n_cholesky );
        formic::xcopy(_norb*_norb, &col.at(0), 1, &_cholesky_vecs.at(_norb*_norb*j), 1);
        formic::xcopy(_norb*_norb, &col.at(0), 1, &_cholesky_vecs_t.at(j), _n_cholesky);
      }

      void to_archive(formic::Archive & ar) const;
      void from_archive(formic::Archive & ar);

      void to_archive(const std::string & filename) const;
      void from_archive(const std::string & filename);

  };

  void orthogonalize_orbs_by_overlap(const int n,
                                     const std::vector<double> & overlap,
                                     std::vector<double> & int1,
                                     std::vector<double> & int2);

  void orthogonalize_orbs_by_overlap(const int n,
                                     const std::vector<std::complex<double> > & overlap,
                                     std::vector<std::complex<double> > & int1,
                                     std::vector<std::complex<double> > & int2);

}

template<class T, class S> formic::BinaryArchiveTemplate<T> & operator<<(formic::BinaryArchiveTemplate<T> & ar, const formic::Hamiltonian<S> & ham);
template<class T, class S> formic::BinaryArchiveTemplate<T> & operator>>(formic::BinaryArchiveTemplate<T> & ar, formic::Hamiltonian<S> & ham);

#endif
